#!/bin/sh
# check cuda version is 12.x

nvcc_output=$(nvcc --version)
if [[ $? -ne 0 ]]; then
  echo "Error: Failed to execute 'nvcc --version'." >&2
fi

cuda_line=$(echo "$nvcc_output" | grep 'Cuda compilation tools')

if [[ $cuda_line =~ release\ ([0-9]+\.[0-9]+) ]]; then
  cuda_version=${BASH_REMATCH[1]}
  echo "CUDA Version: $cuda_version"
else
  echo "Error: Unable to parse CUDA version from output." >&2
fi

if [[ $nvcc_output =~ release\ ([0-9]+)\.([0-9]+) ]]; then
  cuda_major=${BASH_REMATCH[1]}
  cuda_minor=${BASH_REMATCH[2]}
  cuda_version="$cuda_major.$cuda_minor"

  if [[ $cuda_major -eq 12 ]]; then
    echo "Found CUDA 12.x"
    export CUDA_MAJOR_VERSION=$cuda_major
  else
    echo "Expected CUDA 12.x, found $cuda_version"
  fi
else
  echo "Error: Failed to parse CUDA version." >&2
  exit 1
fi
